package com.sjf.day02.huffmancode;

import cn.hutool.core.lang.hash.Hash;
import org.jetbrains.annotations.NotNull;

import java.io.*;
import java.util.*;

/**
 * 哈夫曼编码 压缩字符串
 *
 * @author <huangtang>
 * @since 2022/1/10 11:00
 **/
public class HuffmanCode {

    static Map<Byte,String> huffmanCodes = new HashMap<>();
    static StringBuilder sb = new StringBuilder();

    public static void main(String[] args) {
        /*
        String str = "i like like like java do you like a java";
        byte[] contentBytes = str.getBytes();
        byte[] zip = huffmanZip(contentBytes); //获得压缩后的字节数组
        System.out.println("赫夫曼压缩后的字节数组："+Arrays.toString(zip));
        System.out.println("解码后的字符数组为："+ new String(decode(huffmanCodes, zip)));

        */

        //System.out.println("测试压缩文件的代码");
        //String filesrc = "C:\\Users\\sjf\\Desktop\\1.png";
        //String dstfile = "C:\\Users\\sjf\\Desktop\\dst.zip";
        //
        //zipFile(filesrc,dstfile);
        //System.out.println("压缩文件成功");

        System.out.println("解压文件");
        String filesrc = "C:\\Users\\sjf\\Desktop\\dst.zip";
        String dstfile = "C:\\Users\\sjf\\Desktop\\dst.png";
        unzipFile(filesrc,dstfile);
        System.out.println("解压成功");

        //System.out.println(byteToBitString(true,(byte) 28)); // 0000 0010 0000 0001 -> 0

        //System.out.println("未压缩前的长度："+contentBytes.length); //40
        //
        //List<Node> nodes = getNodes(contentBytes);
        //Node node = createHuffmanTree(nodes);
        //
        //System.out.println("前序遍历");
        //node.preOrder();
        //
        //Map<Byte, String> huffmanCodes = getCodes(node);
        //System.out.println("生成的huffman编码表为："+huffmanCodes);
        //
        //byte[] zip = zip(contentBytes, huffmanCodes);
        //System.out.println("huffman编码压缩过后的byte数组："+Arrays.toString(zip)); //17
    }

    /**
     * 压缩文件
     * @param filesrc 传入文件的全路径
     * @param dstFile 压缩文件存放路径
     */
    public static void zipFile(String filesrc, String dstFile){

        FileInputStream is = null;
        //创建输出流和输入流
        OutputStream os = null;
        ObjectOutputStream oos = null;
        try {
             is = new FileInputStream(filesrc);
            //创建一个和源文件大小相同的byte数组
            byte[] b = new byte[is.available()];
            //读取文件
            is.read(b);
            //对文件进行压缩
            byte[] huffmanBytes = huffmanZip(b);
            os = new FileOutputStream(dstFile);
            oos = new ObjectOutputStream(os);

            oos.writeObject(huffmanBytes);
            oos.writeObject(huffmanCodes);
        } catch (Exception e) {
            System.out.println(e.getMessage());
        }finally {
            try {
                is.close();
                os.close();
                oos.close();
            } catch (IOException e) {
                System.out.println(e.getMessage());
            }
        }
    }

    /**
     * 解压文件
     * @param filesrc 解压的文件路径
     * @param dstFile 解压的文件放在哪个位置
     */
    public static void unzipFile(String filesrc, String dstFile){
        FileInputStream is = null;
        ObjectInputStream ois = null;

        OutputStream os = null;

        try {
            is = new FileInputStream(filesrc);
            ois = new ObjectInputStream(is);
            //读取byte数组
            byte[] huffmanBytes = (byte[]) ois.readObject();
            //读取赫夫曼编码表
            Map<Byte,String> codes = (Map<Byte, String>) ois.readObject();
            //解码
            byte[] bytes = decode(codes, huffmanBytes);
            //将解码得到的字节数组写入到目标文件
            os = new FileOutputStream(dstFile);
            os.write(bytes);
        } catch (Exception e) {
            System.out.println(e.getMessage());
        } finally {
            try {
                os.close();
                is.close();
                ois.close();
            } catch (IOException e) {
                System.out.println(e.getMessage());
            }
        }
    }

    /**
     * 解码哈夫曼编码字节数组
     * @param huffmanCodes 哈夫曼编码表
     * @param huffmanBytes 哈夫曼编码后的字节数组
     * @return 初始字符串对应的字符数组
     */
    private static byte[] decode(Map<Byte,String> huffmanCodes,byte[] huffmanBytes){
        StringBuilder sb = new StringBuilder();
        //得到huffmanBytes对应的二进制字符串
        for (int i = 0; i < huffmanBytes.length; i++) {
            boolean flag = (i == huffmanBytes.length -1);
            sb.append(byteToBitString(!flag,huffmanBytes[i]));
        }
        //System.out.println("sb:"+sb);
        HashMap<String, Byte> map = new HashMap<>();
        for(Map.Entry<Byte,String> entry : huffmanCodes.entrySet()){
            map.put(entry.getValue(),entry.getKey());
        }
        List<Byte> list = new ArrayList<>();
        for (int i = 0; i < sb.length();) {
            int count = 1;
            boolean flag = true;
            Byte b = null;
            while(flag){
                String str = sb.substring(i,i+count); //通过count向后扫描字符串，知道扫描到
                b = map.get(str); //获取每次扫描的value
                if(b!=null){
                    flag = false; //如果不为空，说明扫描到了，退出循环
                }else{
                    count++; //为空则继续向后扫描
                }
            }
            list.add(b);
            i+=count; //让i直接移动到count位置
        }
        byte[] bytes = new byte[list.size()];
        for (int i = 0; i < list.size(); i++) {
            bytes[i] = list.get(i);
        }

        return bytes;
    }

    /** 将byte转成一个二进制的字符串
     * @param b 字节 <br>
     * @param flag 表示是否需要高位补0
     *
     * true 需要补0
     * false 不需要补0
     *
     * @return byte对应的二进制字符串
     */
    private static String byteToBitString(boolean flag, byte b){
        int temp = b;
        if(flag){
            temp |= 256;
        }
        String str = Integer.toBinaryString(temp);
        if (flag) {
            return str.substring(str.length() - 8);
        }else{
            return str;
        }
        //System.out.println("str "+ str);
    }


    /**
     * 将字符数组进行哈夫曼编码
     * @param bytes 需要压缩的字符数组
     * @return 返回编码后的字符数组
     */
    private static byte[] huffmanZip(byte[] bytes){
        List<Node> nodes = getNodes(bytes);
        //创建赫夫曼树
        Node node = createHuffmanTree(nodes);
        //生成对应的赫夫曼编码
        Map<Byte, String> huffmanCodes = getCodes(node);
        //对字节数组进行压缩
        byte[] zip = zip(bytes, huffmanCodes);
        //返回压缩后的字节数组
        return zip;
    }

    /**
     * 根据字符数组 获得对应 Node集合
     * @param bytes 字符数组
     * @return Node集合
     */
    private static List<Node> getNodes(byte[] bytes){
        ArrayList<Node> nodes = new ArrayList<>();
        HashMap<Byte, Integer> map = new HashMap<>();
        //统计每个字符出现的次数
        for (byte b : bytes) {
            Integer count = map.get(b);
            if (count == null) {
                map.put(b, 1);
            } else {
                map.put(b, count+1);
            }
            //map.get(b) == 0 ? map.put(b,1) : map.put(b,map.get(b)+1);
        }
        //把每一个键值对转成壹个Node对象，并加入到nodes集合中
        for(Map.Entry<Byte,Integer> entry:map.entrySet()){
            nodes.add(new Node(entry.getKey(),entry.getValue()));
        }

        return nodes;
    }

    /**
     * 根据Node集合构建HuffmanTree
     * @param nodes 头节点
     * @return 返回哈夫曼树的头节点
     */
    private static Node createHuffmanTree(List<Node> nodes){
        while(nodes.size()>1){
            Collections.sort(nodes);

            Node leftNode = nodes.get(0);
            Node rightNode = nodes.get(1);

            Node parent = new Node(null,leftNode.weight + rightNode.weight);

            parent.left = leftNode;
            parent.right = rightNode;

            nodes.remove(leftNode);
            nodes.remove(rightNode);

            nodes.add(parent);
        }
        return nodes.get(0); //返回nodes的根节点，即哈夫曼树的根节点
    }

    private static Map<Byte,String> getCodes(Node root){
        if(root == null){
            return null;
        }
        getCodes(root.left,"0",sb);
        getCodes(root.right,"1",sb);
        return huffmanCodes;
    }

    /**
     * 根据传入的node节点生成对应的huffmanTree编码表
     * @param node node节点
     * @param code 对应的编码
     * @param sb 字符串拼接编码路径
     */
    private static void getCodes(Node node,String code,StringBuilder sb){
        StringBuilder sb2 = new StringBuilder(sb);
        sb2.append(code);
        if(node != null){ //如果node == null 不处理
            //判断当前节点是叶子节点还是非叶子节点
            if(node.data == null){
                //非叶子节点
                //递归处理
                getCodes(node.left,"0",sb2);
                getCodes(node.right,"1",sb2);
            }else{
                //叶子节点
                huffmanCodes.put(node.data,sb2.toString());
            }

        }
    }

    /**
     * 根据传入的字符数组和哈夫曼编码表压缩字符
     * @param bytes 字符数组
     * @param huffmanCodes 哈夫曼编码表
     * @return 返回编码后的字符数组
     */
    private static byte[] zip(byte[] bytes,Map<Byte,String> huffmanCodes){
        StringBuilder stringBuilder = new StringBuilder();
        for (byte b : bytes) {
            stringBuilder.append(huffmanCodes.get(b));
        }
        int len = stringBuilder.length() % 8 == 0 ?
                stringBuilder.length()/8:
                stringBuilder.length()/8 + 1;

        byte[] huffmanCodeBytes = new byte[len];
        for (int i = 0; i < stringBuilder.length(); i+=8) {
            String strByte;
            if(i+8>stringBuilder.length()){
                //最后一次不够8位，不能截取
                strByte = stringBuilder.substring(i);
            }else{
                strByte = stringBuilder.substring(i,i+8);
            }
            huffmanCodeBytes[i/8] = (byte) Integer.parseInt(strByte,2);
        }

        return huffmanCodeBytes;
    }
}

class Node implements Comparable<Node>{
    Byte data; //存放数据 a = '97'
    int weight; //权值，表示字符出现的次数
    Node left;
    Node right;

    public Node(Byte data, int weight) {
        this.data = data;
        this.weight = weight;
    }

    public void preOrder(){
        System.out.println(this);
        if(this.left!=null){
            this.left.preOrder();
        }
        if(this.right!=null){
            this.right.preOrder();
        }
    }

    @Override
    public int compareTo(@NotNull Node o) {
        return this.weight - o.weight; //根据权值从小到大排序
    }

    @Override
    public String toString() {
        return "Node{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }
}
